from sqlalchemy import update

from model.db import connect, reservoir_sensitivity
from model.db.well_model import TbDvpunitbaisc, TbReservoirunitbaisc, TbWellbasic, TbRockCore, TbCoreData, TbCoreParam
from util.exception_handlers import DataNotDelete


def get_dev_unit():
    """
    获取开发单元
    """
    unit = connect.projSession.query(TbDvpunitbaisc).all()
    return unit


def update_dev_unit(add, delete, update):
    """
    更新开发单元
    """
    try:
        if len(add) != 0:
            for item in add:
                newDev = TbDvpunitbaisc(
                    DvpUnit_Name=item['DvpUnit_Name'],
                    Belong_DvpUnit_Name=item['Belong_DvpUnit_Name'],
                    Remark=item['Remark']
                )
                if item['Belong_DvpUnit_Name'] != "":
                    parent = connect.projSession.query(TbDvpunitbaisc).filter(
                        TbDvpunitbaisc.DvpUnit_Name == item['Belong_DvpUnit_Name']).first()
                    if parent is not None:
                        newDev.DvpUnit_Rank = int(parent.DvpUnit_Rank) + 1
                else:
                    newDev.DvpUnit_Rank = 0
                connect.projSession.add(newDev)
                connect.projSession.commit()
                item['DvpUnit_ID'] = newDev.DvpUnit_ID
                item['DvpUnit_Rank'] = newDev.DvpUnit_Rank
        if len(delete) != 0:
            for item in delete:
                connect.projSession.query(TbDvpunitbaisc).filter(TbDvpunitbaisc.DvpUnit_ID == item['DvpUnit_ID']) \
                    .delete(synchronize_session=False)

        if len(update) != 0:
            for item in update:
                connect.projSession.query(TbDvpunitbaisc).filter(TbDvpunitbaisc.DvpUnit_ID == item['DvpUnit_ID']) \
                    .update(item)
        connect.projSession.commit()
        connect.projSession.flush()
        return add
    except Exception as e:
        connect.projSession.rollback()
        raise e


def get_reservoir_by_devunits(devs):
    """
        获取储层单元
    """
    if len(devs) != 0:
        minUnit = [item[-1] for item in devs]
        reservoir_unit = connect.projSession.query(TbReservoirunitbaisc).filter(
            TbReservoirunitbaisc.Belong_Dev_unit.in_(minUnit)).all()
    else:
        reservoir_unit = connect.projSession.query(TbReservoirunitbaisc).all()
    return reservoir_unit


def get_reservoir_list():
    reservoir_units = connect.projSession.query(TbReservoirunitbaisc).all()
    return reservoir_units


def update_reservoir_unit(add, delete, update):
    """
    更新储层单元
    """
    try:
        if len(add) != 0:
            for item in add:
                newReservoir = TbReservoirunitbaisc(**item)
                connect.projSession.add(newReservoir)
                connect.projSession.commit()
                item['Reservoir_ID'] = newReservoir.Reservoir_ID
        if len(delete) != 0:
            for item in delete:
                connect.projSession.query(TbReservoirunitbaisc).filter(
                    TbReservoirunitbaisc.Reservoir_ID == item['Reservoir_ID']) \
                    .delete(synchronize_session=False)

        if len(update) != 0:
            for item in update:
                connect.projSession.query(TbReservoirunitbaisc).filter(
                    TbReservoirunitbaisc.Reservoir_ID == item['Reservoir_ID']) \
                    .update(item)
        connect.projSession.commit()
        connect.projSession.flush()
        return add
    except Exception as e:
        connect.projSession.rollback()
        raise e


def get_well_by_reservoir(rsv):
    """
        获取井单元
    """
    if len(rsv) != 0:
        well_unit = connect.projSession.query(TbWellbasic).filter(TbWellbasic.Belong_Unit_Name.in_(rsv)).all()
    else:
        well_unit = connect.projSession.query(TbWellbasic).all()
    return well_unit


def get_well_list():
    """
        获取井单元
    """
    well_unit = connect.projSession.query(TbWellbasic).all()
    return well_unit


def update_pit_unit(add, delete, update):
    """
    更新井
    """
    try:
        newWellList = []
        if len(add) != 0:
            for item in add:
                newWell = TbWellbasic(**item)
                connect.projSession.add(newWell)
                connect.projSession.commit()
                newWellList.append(newWell)
        if len(delete) != 0:
            for item in delete:
                connect.projSession.query(TbWellbasic).filter(
                    TbWellbasic.Well_ID == item['Well_ID']) \
                    .delete(synchronize_session=False)

        if len(update) != 0:
            for item in update:
                connect.projSession.query(TbWellbasic).filter(
                    TbWellbasic.Well_ID == item['Well_ID']) \
                    .update(item)
        connect.projSession.commit()
        connect.projSession.flush()
        return newWellList
    except Exception as e:
        connect.projSession.rollback()
        raise e


def get_core_param():
    return connect.projSession.query(TbCoreParam).all()


def update_core_param(add, delete, update):
    """
    更新岩芯参数
    """
    try:
        if len(add) != 0:
            for item in add:
                newParam = TbCoreParam(**item)
                connect.projSession.add(newParam)
                connect.projSession.commit()
                item['Parameter_ID'] = newParam.Parameter_ID
        if len(delete) != 0:
            for item in delete:
                delete_item = connect.projSession.query(TbCoreParam).filter(
                    TbCoreParam.Parameter_Ecode == item['Parameter_Ecode']).first()
                if delete_item.Sensitivity_Type is not None:
                    raise DataNotDelete(delete_item.Parameter_Name)
                connect.projSession.delete(delete_item)

        if len(update) != 0:
            for item in update:
                connect.projSession.query(TbCoreParam).filter(
                    TbCoreParam.Parameter_ID == item['Parameter_ID']) \
                    .update(item)
        connect.projSession.commit()
        connect.projSession.flush()
        return add
    except Exception as e:
        connect.projSession.rollback()
        raise e


def update_core_dataList(add, delete, update):
    """
    更新岩芯数据
    """
    try:
        if len(add) != 0:
            for data in add:

                if 'Core_ID' in data and data['Core_ID'] != '':
                    rock_core = connect.projSession.query(TbRockCore).filter_by(Core_ID=data['Core_ID']).first()
                else:
                    rock_core = TbRockCore(Core_Name=data['Core_Name'])
                    connect.projSession.add(rock_core)
                    connect.projSession.commit()
                    data['Core_ID'] = rock_core.Core_ID
                for key, value in data.items():
                    if key not in ['Core_ID', 'Core_Name']:
                        core_param = connect.projSession.query(TbCoreParam).filter_by(Parameter_Ecode=key).first()
                        if value == '':
                            value = 0
                        core_data = TbCoreData(Core_ID=rock_core.Core_ID, Parameter_ID=core_param.Parameter_ID,
                                               Parameter_Value=value)
                        connect.projSession.add(core_data)

        if len(delete) != 0:
            for data in delete:
                connect.projSession.query(TbCoreData).filter_by(Core_ID=data['Core_ID']).delete(
                    synchronize_session=False)

        if len(update) != 0:
            for data in update:
                # Check if the TbRockCore record exists
                rock_core = connect.projSession.query(TbRockCore).filter_by(Core_ID=data['Core_ID']).first()
                # if not rock_core:
                #     rock_core = TbRockCore(Core_ID=data['Core_ID'], Core_Name=data['Core_Name'])
                #     connect.projSession.add(rock_core)

                # Check and create TbCoreParam records
                for key, value in data.items():
                    if key not in ['Core_ID', 'Core_Name', 'Model_ID']:
                        core_param = connect.projSession.query(TbCoreParam).filter_by(Parameter_Ecode=key).first()
                        if not core_param:
                            core_param = TbCoreParam(Parameter_Ecode=key)
                            connect.projSession.add(core_param)

                        # Check and update TbCoreData records
                        core_data = (
                            connect.projSession.query(TbCoreData)
                            .filter_by(Core_ID=rock_core.Core_ID, Parameter_ID=core_param.Parameter_ID)
                            .first()
                        )
                        if not core_data:
                            core_data = TbCoreData(Core_ID=rock_core.Core_ID, Parameter_ID=core_param.Parameter_ID)
                            connect.projSession.add(core_data)

                        core_data.Parameter_Value = value
        connect.projSession.commit()
        connect.projSession.flush()
        return add
    except Exception as e:
        connect.projSession.rollback()
        raise e


def update_core_basic(add, delete, update):
    """
    更新岩芯
    """
    try:
        if len(add) != 0:
            for item in add:
                newCore = TbRockCore(**item)
                connect.projSession.add(newCore)
                connect.projSession.commit()
                item['Core_ID'] = newCore.Core_ID
        if len(delete) != 0:
            for item in delete:
                # 这种删除写法，才能触发事件
                core = connect.projSession.query(TbRockCore).filter(
                    TbRockCore.Core_ID == item['Core_ID']).first()
                connect.projSession.delete(core)

        if len(update) != 0:
            for item in update:
                connect.projSession.query(TbRockCore).filter(
                    TbRockCore.Core_ID == item['Core_ID']) \
                    .update(item)
        connect.projSession.commit()
        connect.projSession.flush()
        return add
    except Exception as e:
        connect.projSession.rollback()
        raise e


def get_core():
    """
        获取岩芯
    """
    core = connect.projSession.query(TbRockCore).all()
    return core


def get_core_by_well(wells):
    """
        根据井获取岩芯
    """
    if len(wells) == 0:
        core = connect.projSession.query(TbRockCore).all()
    else:
        core = connect.projSession.query(TbRockCore).filter(TbRockCore.Belong_Well_Name.in_(wells)).all()
    return core


def get_core_dataList():
    """
        获取岩芯数据
    """
    data = connect.projSession.query(
        TbRockCore.Core_Name,
        TbCoreData.Core_ID,
        TbCoreParam.Parameter_Ecode,
        TbCoreData.Parameter_Value
    ) \
        .join(TbRockCore, TbRockCore.Core_ID == TbCoreData.Core_ID) \
        .join(TbCoreParam, TbCoreParam.Parameter_ID == TbCoreData.Parameter_ID) \
        .order_by(TbCoreParam.Parameter_ID) \
        .all()

    return data


def get_core_dataList_by_pit(pits):
    """
        获取岩芯数据
    """
    if len(pits) != 0:
        data = connect.projSession.query(
            TbRockCore.Core_Name,
            TbCoreData.Core_ID,
            TbCoreParam.Parameter_Ecode,
            TbCoreData.Parameter_Value
        ) \
            .join(TbRockCore, TbRockCore.Core_ID == TbCoreData.Core_ID) \
            .join(TbCoreParam, TbCoreParam.Parameter_ID == TbCoreData.Parameter_ID) \
            .filter(TbRockCore.Belong_Well_Name.in_(pits)) \
            .order_by(TbCoreParam.Parameter_ID) \
            .all()
    else:
        data = connect.projSession.query(
            TbRockCore.Core_Name,
            TbCoreData.Core_ID,
            TbCoreParam.Parameter_Ecode,
            TbCoreData.Parameter_Value
        ) \
            .join(TbRockCore, TbRockCore.Core_ID == TbCoreData.Core_ID) \
            .join(TbCoreParam, TbCoreParam.Parameter_ID == TbCoreData.Parameter_ID) \
            .order_by(TbCoreParam.Parameter_ID) \
            .all()
    return data


def update_core_data(Core_ID, Parameter_Ecode, value):
    """
        更新岩芯参数数据
    """
    core_data_record = connect.projSession.query(TbCoreData).filter_by(Core_ID=Core_ID) \
        .join(TbCoreParam, TbCoreData.Parameter_ID == TbCoreParam.Parameter_ID). \
        filter_by(Parameter_Ecode=Parameter_Ecode).first()
    if core_data_record:
        core_data_record.Parameter_Value = value
        connect.projSession.commit()


def init_core_param(session=None):
    if session is None:
        connect_session = connect.projSession
    else:
        connect_session = session
    Sensitivity_Type_list = [0, 1, 2, 3, 4, 5, 6]
    Parameter_Ecode_list = ['index', 'block', 'salt', 'acid', 'alkaline', 'speed', 'pressure']
    Parameter_Name_list = ['水敏', '水锁', '盐敏', '酸敏', '碱敏', '速敏', '压敏']
    for i in Sensitivity_Type_list:
        # core_param = reservoir_sensitivity.query(TbCoreParam).filter_by(Sensitivity_Type=i)
        core_param = connect_session.query(TbCoreParam).filter_by(Sensitivity_Type=i).first()
        if core_param is None:
            new_pram = TbCoreParam(
                Parameter_Ecode=Parameter_Ecode_list[i],
                Parameter_Name=Parameter_Name_list[i],
                Sensitivity_Type=i,
            )
            connect_session.add(new_pram)
            connect_session.commit()
            connect_session.flush(new_pram)
